from PIL import Image, ImageDraw, ImageFont
import os
import platform

# -------------------------- 配置参数（可根据需求修改） --------------------------
# 字体路径：自动适配Windows/macOS/Linux系统，若自定义路径需替换下方对应值
SYSTEM = platform.system()
if SYSTEM == "Windows":
    # Windows默认字体路径（可替换为自己的字体路径，如"C:/Windows/Fonts/simhei.ttf"）
    FONT_PATH = "C:/Windows/Fonts/simhei.ttf"
elif SYSTEM == "Darwin":  # macOS
    # macOS默认字体路径（可替换为"/Library/Fonts/Arial Unicode.ttf"等）
    FONT_PATH = "/Library/Fonts/Songti.ttc"
else:  # Linux（保留原路径，可根据发行版调整）
    FONT_PATH = "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc"

# 输入/输出目录：建议使用绝对路径，避免相对路径出错（例：Windows用"D:/pomeranian/input"）
INPUT_DIR = "../add_text/pomeranian_colors"  # 原始图片目录
OUTPUT_DIR = "../add_text/pomeranian_colors_annotated"  # 加文字后输出目录

# 文本样式配置
TEXT_COLOR = (0, 0, 0)  # 文字颜色（RGB，当前为纯黑）
BACKGROUND_COLOR = (255, 255, 255)  # 文字区域背景色（RGB，当前为纯白）
TEXT_MAX_WIDTH_RATIO = 0.96  # 文字最大宽度占图片宽度的比例（避免文字超出边界）
TEXT_PADDING_RATIO = 0.02  # 文字区域上下左右留白占图片宽度的比例
FONT_SIZE_RATIO = 0.03  # 字体大小占图片宽度的比例（保证不同尺寸图片字体适配）
LINE_SPACING_RATIO = 0.2  # 行间距占单行文字高度的比例

# -------------------------- 25种博美犬毛色数据（与图片一一对应） --------------------------
# 格式：[序号, 英文名称, 中文描述]，确保与actual_file_names的顺序完全一致
color_data = [
    [1, "Orange", "橙色：颜色范围从深锈红色到浓郁的金色，再到浅蜂蜜金色。"],
    [2, "Cream", "奶油色：一种极其浅淡的橙色。"],
    [3, "Orange Sable", "橙色貂色：身体为橙色，底毛为更浅的橙色或奶油色，针毛尖端为黑色。"],
    [4, "Black and Tan", "黑褐色：基底为黑色，在特定部位（如眉毛、口吻、胸部、腿）有褐色斑块。"],
    [5, "Black", "黑色：底毛和针毛均为纯黑色。"],
    [6, "White", "白色：纯白色，不应有柠檬色或奶油色的阴影。"],
    [7, "Wolf Sable", "狼貂色：底毛为灰色，针毛为灰色带黑色尖端。"],
    [8, "Parti Color", "花色：白色为底色，带有其他颜色的斑块，头部有白色条纹为佳。"],
    [9, "Chocolate", "巧克力色：浓郁的棕色。"],
    [10, "Cream Sable", "奶油貂色：极浅的奶油色底毛，带有貂色（通常为黑色或深色）尖端。"],
    [11, "Blue Merle", "蓝色陨石色：黑色基底上带有灰色、蓝色斑块，像大理石花纹。"],
    [12, "Red", "红色：深锈红色。"],
    [13, "Chocolate and Tan", "巧克力褐：基底为巧克力色，带有褐色斑块。"],
    [14, "Red Sable", "红色貂色：深锈红色的底毛，带有貂色（黑色）尖端。"],
    [15, "Tri Colored", "三色：是花色博美，同时身上还带有褐色的斑块图案。"],
    [16, "Blue", "蓝色：是黑色的稀释色，呈现灰蓝色，鼻子、眼眶等应为蓝色。"],
    [17, "Blue Sable", "蓝色貂色：基底为蓝色，针毛为深蓝色，末端为黑色尖端。"],
    [18, "Blue and Tan", "蓝褐：基底为蓝色，带有褐色斑块。"],
    [19, "Chocolate Merle", "巧克力陨石色：巧克力色基底上带有浅棕色斑块，形成大理石花纹。"],
    [20, "Beaver", "海狸色/河狸色：颜色从浅饼干色到深棕灰色不等，毛发表面有银蓝色光泽。"],
    [21, "Black and Brindle", "黑虎斑：底色为金色、红色或橙色，带有强烈的黑色横条纹。"],
    [22, "Chocolate Sable", "巧克力貂色：基底为巧克力色，针毛为更深的巧克力色。"],
    [23, "Lavender", "薰衣草色：一种灰白色，毛发表面带有粉红-巧克力色调。"],
    [24, "Blue Brindle", "蓝虎斑：底色为金色、红色或橙色，带有强烈的蓝色横条纹。"],
    [25, "Merle and Tan", "陨石褐：在蓝色陨石色或巧克力陨石色的基础上，还带有褐色斑块。"]
]

# 图片文件名列表（与color_data顺序一一对应，确保文件名与原始图片完全一致）
actual_file_names = [
    "01_Orange.png", "02_Cream.png", "03_Orange_Sable.png", "04_Black_and_Tan.png", "05_Black.png",
    "06_White.png", "07_Wolf_Sable.png", "08_Parti_Color.png", "09_Chocolate.png", "10_Cream_Sable.png",
    "11_Blue_Merle.png", "12_Red.png", "13_Chocolate_and_Tan.png", "14_Red_Sable.png", "15_Tri_Colored.png",
    "16_Blue.png", "17_Blue_Sable.png", "18_Blue_and_Tan.png", "19_Chocolate_Merle.png", "20_Beaver.png",
    "21_Black_and_Brindle.png", "22_Chocolate_Sable.png", "23_Lavender.png", "24_Blue_Brindle.png",
    "25_Merle_and_Tan.png"
]


# -------------------------- 工具函数 --------------------------
def get_text_size(draw, text, font):
    """
    获取文本的宽度和高度
    :param draw: ImageDraw对象
    :param text: 待计算的文本
    :param font: 字体对象
    :return: (width, height) 文本宽高
    """
    bbox = draw.textbbox((0, 0), text, font=font)  # 获取文本边界框（left, top, right, bottom）
    return (bbox[2] - bbox[0], bbox[3] - bbox[1])


def wrap_text(draw, text, font, max_width):
    """
    按最大宽度自动换行（适配中文，按字符分割）
    :param draw: ImageDraw对象
    :param text: 待换行的文本
    :param font: 字体对象
    :param max_width: 文本最大宽度
    :return: 换行后的文本列表（每行为一个元素）
    """
    if not text.strip():
        return []

    lines = []
    current_line = ""
    # 中文按单个字符分割，避免词语拆分；英文可保留空格（此处兼容中英文）
    chars = list(text) if all('\u4e00' <= c <= '\u9fff' for c in text.strip()) else text.split()

    for char in chars:
        # 检查当前行加新字符后的宽度（英文需补空格）
        test_line = f"{current_line} {char}" if (not all(
            '\u4e00' <= c <= '\u9fff' for c in char)) and current_line else current_line + char
        test_width, _ = get_text_size(draw, test_line, font)

        if test_width <= max_width:
            current_line = test_line
        else:
            if current_line:  # 避免空行
                lines.append(current_line.strip())
            current_line = char.strip()

    # 添加最后一行
    if current_line.strip():
        lines.append(current_line.strip())

    return lines


def add_text_to_image(input_path, output_path, text):
    """
    给图片添加文字说明（图片顶部保留原图，底部添加白色背景+黑色文字）
    :param input_path: 原始图片路径
    :param output_path: 输出图片路径
    :param text: 要添加的文字说明
    """
    try:
        # 1. 打开原始图片（转为RGB模式，避免透明通道问题）
        with Image.open(input_path).convert("RGB") as img:
            img_width, img_height = img.size
            print(f"处理图片：{os.path.basename(input_path)} | 原始尺寸：{img_width}x{img_height}")

            # 2. 初始化字体（按图片宽度比例设置字体大小，适配不同尺寸图片）
            font_size = int(img_width * FONT_SIZE_RATIO)
            # 容错：若字体路径不存在，使用默认字体
            try:
                font = ImageFont.truetype(FONT_PATH, font_size)
            except IOError:
                font = ImageFont.load_default()
                print(f"警告：未找到指定字体 {FONT_PATH}，使用默认字体")

            # 3. 计算文本换行和所需高度
            draw_temp = ImageDraw.Draw(img)  # 临时绘图对象用于计算文本
            text_max_width = int(img_width * TEXT_MAX_WIDTH_RATIO)  # 文本最大宽度
            wrapped_lines = wrap_text(draw_temp, text, font, text_max_width)  # 自动换行
            line_width, line_height = get_text_size(draw_temp, "测试文本", font)  # 单行文本高度
            total_line_spacing = (len(wrapped_lines) - 1) * int(line_height * LINE_SPACING_RATIO)  # 总行间距
            text_area_height = len(wrapped_lines) * line_height + total_line_spacing  # 文本区域高度
            text_area_height += int(img_width * TEXT_PADDING_RATIO * 2)  # 上下留白（按图片宽度比例，避免过窄/过宽）

            # 4. 创建新图片（原图高度 + 文本区域高度，背景为白色）
            new_img_height = img_height + text_area_height
            new_img = Image.new('RGB', (img_width, new_img_height), color=BACKGROUND_COLOR)
            # 将原始图片粘贴到新图片顶部（位置：(0, 0)）
            new_img.paste(img, (0, 0))

            # 5. 在新图片底部绘制文本
            draw = ImageDraw.Draw(new_img)
            text_x = int(img_width * TEXT_PADDING_RATIO)  # 文本左对齐，留左留白
            text_y = img_height + int(img_width * TEXT_PADDING_RATIO)  # 文本顶部在原图底部+上留白

            # 逐行绘制文本
            for line in wrapped_lines:
                draw.text((text_x, text_y), line, font=font, fill=TEXT_COLOR)
                text_y += line_height + int(line_height * LINE_SPACING_RATIO)  # 下移到下一行

            # 6. 保存图片（保留PNG格式，确保清晰度）
            new_img.save(output_path, format="PNG", quality=95)
            print(f"✅ 成功保存：{os.path.basename(output_path)} | 最终尺寸：{img_width}x{new_img_height}\n")

    except Exception as e:
        print(f"❌ 处理 {os.path.basename(input_path)} 失败：{str(e)}\n")


# -------------------------- 主程序（批量处理） --------------------------
if __name__ == "__main__":
    # 1. 检查输入目录是否存在
    if not os.path.exists(INPUT_DIR):
        print(f"错误：输入目录 {INPUT_DIR} 不存在，请检查路径是否正确！")
        exit(1)

    # 2. 创建输出目录（若不存在）
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    print(f"输出目录已准备：{OUTPUT_DIR}\n")

    # 3. 批量处理每张图片（确保color_data、actual_file_names顺序一致）
    if len(color_data) != len(actual_file_names):
        print(f"错误：毛色数据（{len(color_data)}条）与图片文件名（{len(actual_file_names)}个）数量不匹配！")
        exit(1)

    # 遍历所有图片和对应数据
    for idx, (data, filename) in enumerate(zip(color_data, actual_file_names)):
        color_num, color_en, color_cn = data
        input_path = os.path.join(INPUT_DIR, filename)

        # 检查原始图片是否存在
        if not os.path.exists(input_path):
            print(f"⚠️跳过：{filename} 未在输入目录找到，路径：{input_path}\n")
            continue

        # 构造要添加的文字（格式：序号 中文名称（英文）+ 描述）
        text_to_add = f"{str(color_num).zfill(2)} {color_cn.split('：')[0]}（{color_en}）：{color_cn.split('：')[1]}"
        # 输出路径（与输入文件名一致，便于对应）
        output_path = os.path.join(OUTPUT_DIR, filename)

        # 调用函数添加文字
        add_text_to_image(input_path, output_path, text_to_add)

    print("=" * 50)
    print("批量处理完成！")
    print(f"原始图片目录：{os.path.abspath(INPUT_DIR)}")
    print(f"处理后图片目录：{os.path.abspath(OUTPUT_DIR)}")
    print("=" * 50)
